import torch
import torch.nn as nn

tagset_size = 5
START_TAG = "<START>"
STOP_TAG = "<STOP>"
tag_to_ix = {"B": 0, "I": 1, "O": 2, START_TAG: 3, STOP_TAG: 4}
transitions = nn.Parameter(torch.randn(tagset_size,tagset_size))


def argmax(vec):
    _, idx = torch.max(vec, 1) #返回每行的最大的元素和索引


def _viterbi_decode(feats):
    # 预测序列的得分，维特比解码，输出得分与路径值
    backpointers = []

    #初始化
    init_vvars = torch.full((1, tagset_size), -1000.)
    init_vvars[0][tag_to_ix[START_TAG]] = 0 #这就保证了一定是从START到其他标签,只有在START_TAG的位置是0

    # forward_var at step i holds the viterbi variables for step i-1

    forward_var = init_vvars

    for feat in feats:
        bptrs_t = []  # holds the backpointers for this step
        viterbivars_t = []  # holds the viterbi variables for this step

        for next_tag in range(tagset_size):
        # 其他标签（B,I,E,Start,End）到标签next_tag的概率

            next_tag_var = next_tag

    return init_vvars


if __name__ == '__main__':
    print(_viterbi_decode(1))